import time
import torch
import ghhops_server as hs
import rhino3dm as r3d
import json

from flask import Flask
from tqdm import tqdm

app = Flask(__name__)
hops = hs.Hops(app)


def distance_point_to_line_segment(px, py, ax, ay, bx, by):
    ## px, py (n,) ax,bx,ay,by (k,)

    dx1, dy1 = (bx - ax).view(-1,1), (by - ay).view(-1,1)
    # dx1 (k,1) dy1 (k,1)
    # AP
    dx2, dy2 = px.view(1, -1) - ax.view(-1,1), py.view(1,-1) - ay.view(-1,1)
    # dx2 (k,n) dy2 (k,n)


    length_squared = dx1 ** 2 + dy1 ** 2
    # length_squared (k)
    # # 线段长度为零
    # if length_squared == 0:
    #     return math.sqrt(dx2 ** 2 + dy2 ** 2)

    # 投影标量t
    t = torch.max(torch.tensor(0), torch.min(torch.tensor(1), (dx2 * dx1 + dy2 * dy1) / length_squared))

    # t (k,n)

    # 投影坐标
    projx = ax.reshape(-1,1) + t * dx1
    projy = ay.reshape(-1,1) + t * dy1

    # P到投影点距离
    return torch.sqrt((px.reshape(1,-1) - projx) ** 2 + (py.reshape(1,-1) - projy) ** 2) # (k,n)


def rotate_point_2d(point: torch.Tensor,
                    angle: torch.Tensor,
                    center: torch.Tensor) -> torch.Tensor:
    """
    点绕中心点旋转（2D, XY 平面）

    参数:
        point:  shape=(2,) or (3,), 要旋转的点
        angle:  标量 Tensor，旋转弧度
        center: shape=(2,) or (3,), 旋转中心

    返回:
        rotated: 和输入同 shape 的 Tensor
    """

    # 确保是 float tensor
    point = point.to(dtype=torch.float32)
    center = center.to(dtype=torch.float32)

    # 转到中心坐标系
    p = point[..., :2] - center[..., :2]

    # 构造旋转矩阵
    cos_a = torch.cos(angle)
    sin_a = torch.sin(angle)
    R = torch.stack([
        torch.stack([cos_a, -sin_a]),
        torch.stack([sin_a,  cos_a])
    ])  # (2,2)

    # 应用旋转
    p_rot = R @ p

    # 回到原坐标系
    result_xy = p_rot + center[..., :2]

    # 如果是 3D 点，把 Z 维拼回去
    if point.shape[0] == 3:
        return torch.cat([result_xy, point[2:].clone()])
    else:
        return result_xy


def point_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    计算两点欧几里得距离

    参数:
        a: Tensor, shape=(d,) 或 (N,d)
        b: Tensor, shape=(d,) 或 (N,d)

    返回:
        dist: 标量 (d,) 情况 或 shape=(N,) 的 Tensor
    """
    return torch.norm(a - b, dim=-1)


@hops.component(
    "/prototype_test",                                  # 端点路径
    name="prototype_test", description="prototype_test",
    inputs=[
        hs.HopsString("init_params"),
        hs.HopsPoint("base_pt"),
        hs.HopsPoint("circle_pt"),
        hs.HopsNumber("circle_radius"),

        hs.HopsPoint("constraint_pt"),
        hs.HopsNumber("constraint_radius"),
    ],
    outputs=[hs.HopsString("params")]
)
def prototype_test(
    init_params: str, base_pt: r3d.Point3d,
    circle_pt: r3d.Point3d, circle_radius: float,
    constraint_pt: r3d.Point3d, constraint_radius: float,
):

    params =  torch.tensor([float(i) for i in init_params.split()], requires_grad=True, dtype=torch.float32)
    base_pt = torch.tensor([base_pt.X, base_pt.Y, base_pt.Z], dtype=torch.float32)
    circle_pt = torch.tensor([circle_pt.X, circle_pt.Y, circle_pt.Z], dtype=torch.float32)
    circle_radius = torch.tensor(float(circle_radius), dtype=torch.float32)
    constraint_pt = torch.tensor([constraint_pt.X, constraint_pt.Y, constraint_pt.Z], dtype=torch.float32)
    constraint_radius = torch.tensor(float(constraint_radius), dtype=torch.float32)
    stick_len = 100.

    # iters: int = 0
    iters: int = 500
    lr: float = 0.05
    tol: float = 1e-6
    opt = torch.optim.Adam([params], lr=lr)

    def forward(params: torch.Tensor) -> torch.Tensor:

        last = base_pt
        distance_to_constraint_circle = torch.tensor([0 for _ in range(len(params))], dtype=torch.float32)
        for idx, param in enumerate(params):
            temp = last + torch.tensor([stick_len, 0, 0])
            temp = rotate_point_2d(temp, param, last)

            temp_distance = distance_point_to_line_segment(
                constraint_pt[0], constraint_pt[1],
                last[0], last[1], temp[0], temp[1]
            )
            distance_to_constraint_circle[idx] = torch.clamp(constraint_radius-temp_distance, min=0.)

            last = temp

        cost = torch.mean(distance_to_constraint_circle)
        print(f"cost: {cost}")

        # the last node distance to the target circle edge
        distance = point_distance(last, circle_pt) - circle_radius

        loss = torch.clamp(distance, min=0.0) + cost**2


        return loss


    history_params = [params.tolist()]

    last_loss = None
    iter_idx = 0
    for _ in tqdm(range(iters)):
        iter_idx += 1
        loss = forward(params)

        opt.zero_grad()
        loss.backward()
        opt.step() # modify params

        history_params.append(
            params.tolist()
        )

        current = loss.item()
        if last_loss is not None and abs(current - last_loss) < tol:
            break

        last_loss = current

    with torch.no_grad():
        final_loss = forward(params)

    print(f"final loss: {final_loss.item()} after {iter_idx} iterations")

    return json.dumps(history_params)



if __name__ == "__main__":
    app.run(debug=True)
